/*
 * jngc, Java utility classes by Nicholas Coltharp
 * Copyright (C) 2010 Nicholas Coltharp <ngc1@rice.edu>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package edu.rice.ngc1.util.thread;

import java.util.Collections;
import java.util.List;
import java.util.Queue;
import java.util.Vector;
import java.util.concurrent.AbstractExecutorService;
import java.util.concurrent.BrokenBarrierException;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.ThreadPoolExecutor.CallerRunsPolicy;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * <p>
 * An {@code ExecutorService} which behaves similarly to the fixed
 * thread pool returned by {@link Executors#newFixedThreadPool(int)}.
 * The main difference is that an object of this class does not queue
 * tasks. Instead, once all of its threads are busy, additional tasks
 * will be run in the thread that submitted them. This is ostensibly
 * the same effects as a {@link ThreadPoolExecutor} with a rejection
 * policy set to {@link CallerRunsPolicy}, but in practice, that has a
 * tendency to deadlock. This class does not deadlock.
 * </p>
 * 
 * <p>
 * If a submitted task throws an unchecked exception, the exception
 * will be allowed to propagate. The thread in which the exception was
 * thrown will be discarded and a new one will take its place. If one
 * of this class's worker threads is interrupted while it is waiting
 * for work, the thread's interrupted status will be set and the
 * thread will be discarded.
 * </p>
 */
public class LimitedThreadsExecutor extends AbstractExecutorService {

    /**
     * The {@code ThreadFactory} which this object will use for thread
     * creation.
     */
    private final ThreadFactory threadFactory;

    /**
     * List of all {@code TaskRunner}s whose threads have not been
     * started.
     */
    private final Queue<TaskRunner> unstartedRunners;

    /**
     * List of all {@code TaskRunner}s whose threads have been started
     * but who are not executing any tasks.
     */
    private final Queue<TaskRunner> availableStartedRunners;

    /**
     * List of all {@code TaskRunner}s owned by this object. Not
     * modified after creation.
     */
    private final List<TaskRunner> allTaskRunners;

    /**
     * Boolean value indicating whether this executor is currently in
     * the process of shutting down.
     */
    private final AtomicBoolean isShuttingDown;

    /**
     * A {@code CountDownLatch} used when this executor is shutting
     * down. All of the task runners activate this latch when they die
     * during shutdown.
     */
    private final CountDownLatch shutDownLatch;

    /**
     * <p>
     * Runnable class which does the actual work of running tasks.
     * </p>
     * 
     * <p>
     * The lifecycle of a {@code TaskRunner} is as follows:
     * <ol>
     * <li>The {@code TaskRunner} is created.</li>
     * <li>The {@code TaskRunner}'s thread is created and assigned.</li>
     * <li>The {@code TaskRunner} is added to the
     * {@link LimitedThreadsExecutor#unstartedRunners} queue.</li>
     * <li>When the {@code TaskRunner} is needed to run a task, its
     * thread is started.</li>
     * <li>The {@code TaskRunner} waits for a thread to wait on its
     * {@link TaskRunner#readyBarrier}.</li>
     * <li>When a thread waits on the {@code readyBarrier}, it is
     * signaling that the {@code TaskRunner}'s task is set and ready
     * to be run. Once the barrier is breached, the task will be run.</li>
     * <li>The thread that starts the runner will remove it from the
     * runner queue.</li>
     * <li>There are a few ways this process can end:
     * <ul>
     * <li>If the task finishes as normal, the runner will re-add
     * itself to the runner queue.</li>
     * <li>If an exception is thrown from within the task, it will be
     * allowed to propagate.</li>
     * <li>If a shutdown begins while the runner is waiting for work,
     * it will be interrupted.</li>
     * <li>If {@link LimitedThreadsExecutor#shutdownNow()} is called
     * while the runner is executing its task, the runner will be
     * interrupted; however, there is no guarantee that it will
     * respond.</li>
     * <li>In any case, when the runner finishes or is interrupted, if
     * a shutdown is in progress, then it will cease its action.
     * Otherwise, if it was not interrupted by an exception, it will
     * wait for more work. If it was interrupted by an exception, it
     * will simply sit until its new thread is started.</li>
     * </ul>
     * </ol>
     * </p>
     */
    private class TaskRunner implements Runnable {

        /**
         * A "slot" where the calling thread can put the task that
         * this runner is supposed to run.
         */
        public Runnable task;

        /**
         * The thread that is running this object.
         */
        public Thread thread;

        /**
         * The thread will hang on this barrier until the calling
         * thread also hangs on it. When this barrier is breached, it
         * indicates that the task has been set.
         */
        public final CyclicBarrier readyBarrier = new CyclicBarrier(2);

        /**
         * Starts this {@code TaskRunner}.
         */
        @Override
        public void run() {
            try {
                while (!isShutdown()) {
                    readyBarrier.await();
                    task.run();
                    availableStartedRunners.add(this);
                }
            } catch (BrokenBarrierException e) {
                readyBarrier.reset();
            } catch (InterruptedException e) {
                if (!isShutdown()) {
                    this.thread.interrupt();
                }
            } finally {
                if (isShutdown()) {
                    shutDownLatch.countDown();
                } else {
                    this.thread = threadFactory.newThread(this);
                    unstartedRunners.add(this);
                }
            }
        }
    }

    /**
     * Constructs a new {@code LimitedThreadsExecutor} with the given
     * number of threads and the {@code ThreadFactory} returned by
     * {@link Executors#defaultThreadFactory()}.
     * 
     * @param nThreads
     *            the number of threads to give this executor
     */
    public LimitedThreadsExecutor(int nThreads) {
        this(nThreads, Executors.defaultThreadFactory());
    }

    /**
     * Constructs a new {@code LimitedThreadsExecutor} with the given
     * number of threads and the given {@code ThreadFactory}.
     * 
     * @param nThreads
     *            the number of threads to give this executor
     * @param threadFactory
     *            the {@code ThreadFactory} to make this executor use
     *            for constructing threads
     */
    public LimitedThreadsExecutor(int nThreads, ThreadFactory threadFactory) {
        if (nThreads < 1) {
            throw new IllegalArgumentException();
        }
        this.shutDownLatch = new CountDownLatch(nThreads);
        this.isShuttingDown = new AtomicBoolean();
        this.threadFactory = threadFactory;
        this.unstartedRunners = new ConcurrentLinkedQueue<TaskRunner>();
        this.availableStartedRunners = new ConcurrentLinkedQueue<TaskRunner>();
        this.allTaskRunners = new Vector<TaskRunner>(nThreads);

        for (int i = 0; i < nThreads; ++i) {
            TaskRunner runner = new TaskRunner();
            Thread thread = threadFactory.newThread(runner);
            runner.thread = thread;
            this.unstartedRunners.add(runner);
            this.allTaskRunners.add(runner);
        }
    }

    /**
     * Starts one thread owned by this object that has not already
     * been started, or does nothing if no such thread exists.
     */
    public void startAThread() {
        if (!isShutdown()) {
            TaskRunner runner = unstartedRunners.poll();
            if (runner != null) {
                runner.thread.start();
                availableStartedRunners.add(runner);
            }
        }
    }

    /**
     * Starts all threads owned by this object that have not already
     * been started.
     */
    public void startAllThreads() {
        if (!isShutdown()) {
            while (!unstartedRunners.isEmpty()) {
                TaskRunner runner = unstartedRunners.remove();
                runner.thread.start();
                availableStartedRunners.add(runner);
            }
        }
    }

    /**
     * Executes the given {@code Runnable}. If there are threads
     * available, the command will be run in a separate thread and
     * this method will immediately return. Otherwise, the command
     * will be run in the current thread.
     * 
     * @param command
     *            the {@code Runnable} to run
     */
    @Override
    public void execute(Runnable command) {
        if (!isShutdown()) {
            TaskRunner taskRunner = availableStartedRunners.poll();
            if (taskRunner == null) {
                taskRunner = unstartedRunners.poll();
                if (taskRunner == null) {
                    command.run();
                    return;
                } else {
                    taskRunner.thread.start();
                }
            }
            taskRunner.task = command;
            try {
                /*
                 * the runner will proceed once both it and this are
                 * waiting on this barrier
                 */
                taskRunner.readyBarrier.await();
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            } catch (BrokenBarrierException e) {
                // the other thread handles it
            }
        }
    }

    /**
     * Stops all of this object's idle threads and blocks it from
     * receiving any more requests.
     */
    @Override
    public void shutdown() {
        if (!isShuttingDown.getAndSet(true)) {
            for (TaskRunner runner : availableStartedRunners) {
                runner.thread.interrupt();
            }
        }
    }

    /**
     * In addition to the effects of {@link #shutdown()}, attempts to
     * stop all currently running tasks. This method does not attempt
     * to stop tasks that were not started in separate threads.
     * 
     * @return {@link Collections#emptyList()}
     */
    @Override
    public List<Runnable> shutdownNow() {
        if (!isShuttingDown.getAndSet(true)) {
            for (TaskRunner runner : allTaskRunners) {
                runner.thread.interrupt();
            }
        }
        return Collections.emptyList();
    }

    /**
     * Indicates whether a shutdown has been requested.
     * 
     * @return {@code true} if a shutdown has been requested
     */
    @Override
    public boolean isShutdown() {
        return isShuttingDown.get();
    }

    /**
     * Waits for all of this object's threads to die or for the given
     * timeout to be reached, whichever comes first.
     * 
     * @param timeout
     *            the amount of time to wait
     * @param unit
     *            the unit in which {@code timeout} is expressed
     * @return {@code true} if this object's threads actually
     *         terminated; {@code false} if the timeout was reached
     * @throws InterruptedException
     *             if the current thread is interrupted while waiting
     */
    @Override
    public boolean awaitTermination(long timeout, TimeUnit unit)
            throws InterruptedException {
        return shutDownLatch.await(timeout, unit);
    }

    /**
     * Indicates whether all of this object's threads have terminated.
     * 
     * @return {@code true} if this object's threads have all
     *         terminated; {@code false} otherwise
     */
    @Override
    public boolean isTerminated() {
        return (shutDownLatch.getCount() == 0);
    }

    /**
     * Shuts this object down when it is garbage-collected.
     */
    @Override
    protected void finalize() throws Throwable {
        super.finalize();
        shutdown();
    }
}
